use data_structure::prelude::*;
use std::cmp::min;

fn strong_comp(
    map: &Vexes<Table<DNode<u64>>, ThreadedAVLTree<u64, u64>>,
    nodes: Vec<u64>,
) -> Vec<Vec<u64>> {
    let mut res: Vec<Vec<u64>> = Vec::new();
    let num_of_vexes = map.size() as usize;
    let mut dfn = vec![0; num_of_vexes];
    let mut low = vec![0; num_of_vexes];
    let mut in_stack = vec![false; num_of_vexes];
    let mut stack_nodes: Stack<List<u64>> = Stack::new();
    let mut stack_edges: Stack<List<(u64, u64)>> = Stack::new();
    for (idx, &n) in nodes.iter().enumerate() {
        if dfn[idx] == 0 {
            stack_nodes.push(n);
            dfn[idx] = stack_nodes.size();
            low[idx] = dfn[idx];
            in_stack[idx] = true;
            stack_edges.push((n, n));
            map.edges(n)
                .into_iter()
                .for_each(|(to, _)| stack_edges.push((n, to)));
            while let Some((from, to)) = stack_edges.pop() {
                let idx_from = nodes.iter().position(|x| *x == from).unwrap();
                let idx_to = nodes.iter().position(|x| *x == to).unwrap();
                if from != to {
                    if dfn[idx_to] == 0 {
                        stack_nodes.push(to);
                        in_stack[idx_to] = true;
                        dfn[idx_to] = stack_nodes.size();
                        low[idx_to] = dfn[idx_to];
                        stack_edges.push((idx_from as u64, idx_to as u64));
                        stack_edges.push((to, to));
                        map.edges(to)
                            .into_iter()
                            .for_each(|(tto, _)| stack_edges.push((to, tto)));
                    } else if in_stack[idx_to] {
                        low[idx_from] = min(low[idx_from], dfn[idx_to]);
                    }
                } else {
                    if low[idx_to] == dfn[idx_to] {
                        let mut set = Vec::<u64>::new();
                        while let Some(node) = stack_nodes.pop() {
                            let idx_node = nodes.iter().position(|x| *x == node).unwrap();
                            set.push(node);
                            in_stack[idx_node] = false;
                            if node == to {
                                break;
                            }
                        }
                        res.push(set);
                    }
                    if let Some((u, v)) = stack_edges.pop() {
                        let (idx_from, idx_to) = (u as usize, v as usize);
                        low[idx_from] = min(low[idx_from], low[idx_to]);
                    }
                }
            }
        }
    }
    res
}

#[test]
fn tarjan() {
    let nodes = (1..=6u64).collect::<Vec<_>>();
    let mut map: Vexes<Table<DNode<u64>>, ThreadedAVLTree<u64, u64>> =
        Vexes::new(nodes.clone(), true);
    map.update_edge(2, 1, 1);
    map.update_edge(2, 4, 1);
    map.update_edge(3, 2, 1);
    map.update_edge(3, 6, 1);
    map.update_edge(4, 3, 1);
    map.update_edge(4, 5, 1);
    map.update_edge(4, 6, 1);
    map.update_edge(5, 1, 1);
    map.update_edge(6, 1, 1);
    map.update_edge(6, 2, 1);
    map.update_edge(6, 5, 1);
    let vecs = strong_comp(&map, nodes);
    assert_eq!(vecs.len(), 3);
    assert_eq!(vecs[0], vec![1]);
    assert_eq!(vecs[1], vec![5]);
    assert_eq!(vecs[2], vec![3, 6, 4, 2]);
}
